import backtrader as bt
import os
import sys
import datetime
import akshare as ak
import pandas as pd
from matplotlib import pyplot as plt
from matplotlib_inline import backend_inline

def use_svg_display():
    """使用svg格式在Jupyter中显示绘图

    Defined in :numref:`sec_calculus`"""
    backend_inline.set_matplotlib_formats('svg')

use_svg_display()

class EtfBalance(bt.Strategy):

    def log(self, txt, dt=None, doprint=False):
        ''' 日志函数，用于统一输出日志格式 '''
        if doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

    # 当订单状态变化时触发
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:  # 接受订单交易，正常情况
            return
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log('已买入, 购入价格 {:.2f}, 金额 {:.2f} ,手续费{:.2f}, {}'.format(order.executed.price, order.executed.value,
                                                                      order.executed.comm, order.data._name), doprint=True)
            elif order.issell():
                self.log('已卖出, 卖出价格 {:.2f}, 金额{:.2f} ,手续费{:.2f}, {}'.format(order.executed.price, order.executed.value,
                                                                     order.executed.comm, order.data._name), doprint=True)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消、保证金不足、金额不足拒绝交易', doprint=True)

    def notify_trade(self, trade):
        """
        交易成果

        Arguments:
            trade {object} -- 交易状态
        """
        if not trade.isclosed:
            return

        # 显示交易的毛利率和净利润
        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
                 (trade.pnl, trade.pnlcomm), doprint=False)

    def __init__(self):
        # 起始默认仓位，留点空余给手续费，下同
        self.etf300 = 0.495
        self.etfGold = 0.495
        self.last_month = 0

    def next(self):
        # 时间序列 转 datetime
        # now_ordinal = int(self.data0_datetime[0])
        now_ordinal = int(self.data0.datetime[0])
        now_date = datetime.date.fromordinal(now_ordinal)
        # print('{}'.format(str(now_date)))
        now_month = now_date.month

        if self.last_month == now_month:
            return
        self.last_month = now_month
        # x个月调整一次
        if self.last_month % 12 != 1:
            return

        allAmount = self.broker.getvalue()
        etf300Amount = self.broker.getvalue(datas=[self.data0])

        if etf300Amount/allAmount > self.etf300:
            # etf300市值超过比例线，先调整etf300
            self.order = self.order_target_percent(data=self.data0, target=self.etf300)
            self.order = self.order_target_percent(data=self.data1, target=self.etfGold)
        else:
            self.order = self.order_target_percent(data=self.data1, target=self.etfGold)
            self.order = self.order_target_percent(data=self.data0, target=self.etf300)



def price_scale(df, scale=1.0):
    df['open'] = df['open'] / scale
    df['high'] = df['high'] / scale
    df['low'] = df['low'] / scale
    df['close'] = df['close'] / scale
    return df

def run_test():
    cerebro = bt.Cerebro()

    df_sh000300 = pd.read_csv('csv/sh000300.csv')
    df_sh000300 = df_sh000300.loc[:, ['date', 'open', 'high', 'low', 'close', 'volume']]  # 选择数据
    df_sh000300 = df_sh000300.set_index(pd.to_datetime(df_sh000300['date'].astype('str'))).sort_index()  # 排序
    df_sh000300 = df_sh000300[df_sh000300['date'] > '2017-01-01']
    df_sh000300 = price_scale(df_sh000300, scale=1000)

    df_sh518880 = pd.read_csv('csv/sh518880.csv')
    df_sh518880 = df_sh518880.loc[:, ['date', 'open', 'high', 'low', 'close', 'volume']]  # 选择数据
    df_sh518880 = df_sh518880.set_index(pd.to_datetime(df_sh518880['date'].astype('str'))).sort_index()  # 排序
    df_sh518880 = df_sh518880[df_sh518880['date'] > '2017-01-01']
    df_sh518880 = price_scale(df_sh518880, scale=1)

    datafeed_sh000300 = bt.feeds.PandasData(dataname=df_sh000300,
                                            fromdate=datetime.datetime(2017, 1, 1),
                                            todate=datetime.datetime(2024, 1, 1))
    cerebro.adddata(datafeed_sh000300, name='sh000300')  # 通过 name 实现数据集与股票的一一对应
    print('读取成功')

    datafeed_sh000905 = bt.feeds.PandasData(dataname=df_sh518880,
                                            fromdate=datetime.datetime(2017, 1, 1),
                                            todate=datetime.datetime(2024, 1, 1))
    cerebro.adddata(datafeed_sh000905, name='sh518880')  # 通过 name 实现数据集与股票的一一对应
    print('读取成功')

    cerebro.addstrategy(EtfBalance)
    # 初始资金 1,000,000
    cerebro.broker.setcash(1000000.0)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0001, percabs=True, stocklike=True, interest=0.0, interest_long=True,
                                 name='sh000300')
    cerebro.broker.setcommission(commission=0.0001, percabs=True, stocklike=True, interest=0.0, interest_long=True,
                                 name='sh518880')
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    cerebro.broker.set_coc(True)
    # 启动回测
    result = cerebro.run()

    print('最终持仓', cerebro.broker.getvalue())

    cerebro.plot()


def prepare_csv():
    # 指数
    # index_stock_info_df = ak.index_stock_info()
    # index_stock_info_df.to_csv('csv/指数.csv', index=None)

    # 沪深300
    sh000300_index_daily_df = ak.stock_zh_index_daily(symbol="sh000300")
    sh000300_index_daily_df.to_csv('csv/sh000300.csv', index=None)

    # 黄金ETF
    sh518880_daily_df = ak.stock_zh_index_daily(symbol="sh518880")
    sh518880_daily_df.to_csv('csv/sh518880.csv', index=None)


if __name__ == '__main__':
    # prepare_csv()
    run_test()
